
import numpy as np
import tensorflow as tf
import pathlib
from tensorflow import keras


def load_data(data_dir,img_width = 28, img_height =28, batch_size = 32):
    data_dir = pathlib.Path(data_dir)
    seed = np.random.randint(0,1000)
    print("seed:",seed)
    train_ds = keras.utils.image_dataset_from_directory(
        data_dir,
        seed=seed,
        color_mode='grayscale',
        validation_split=0.2,
        subset="training",
        image_size=(img_height, img_width),
        batch_size=batch_size)
    
    
    val_ds = tf.keras.utils.image_dataset_from_directory(
        data_dir,
        seed=seed,
        color_mode='grayscale',
        validation_split=0.2,
        subset="validation",
        image_size=(img_height, img_width),
        batch_size=batch_size)
    
    return train_ds,val_ds


class EnableMnist(keras.Model):
    
    def __init__(self, num_classes=10):
        super(EnableMnist, self).__init__()
        self.num_classes = num_classes
        self.conv1 = keras.layers.Conv2D(32, 3, activation='relu')
        self.flatten = keras.layers.Flatten()
        self.d1 = keras.layers.Dense(128, activation='relu')
    
        self.conv2 = keras.layers.Conv2D(64, 3, activation='relu')
        self.d2 = keras.layers.Dense(self.num_classes)
    
    def call(self, input):
        x = self.conv1(input)
        x = self.flatten(x)
        x = self.d1(x)
        return x
        
